import cv2
import numpy as np
import os


# face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalface_alt_tree.xml')
# face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalcatface.xml')
# face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalface_alt.xml')
face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalface_alt2.xml')
recognizer = cv2.face.LBPHFaceRecognizer.create()
recognizer.read('./data/03.yml')
cap = cv2.VideoCapture(0)
# cap = cv2.VideoCapture('assets/video/2.mp4')


def face_detect_demo(src):
    gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
    faces = face_detector.detectMultiScale(gray, 1.1, 5)
    for x, y, w, h in faces:
        id, confidence = recognizer.predict(gray[y: y+h, x: x+w])
        print('标签ID：', id, '置信度：', confidence)
        color = (0, 0, 255) if confidence > 60 else (0, 255, 0)
        cv2.rectangle(src, (x, y), (x + w, y + h), color, thickness=2)
        cv2.putText(src, 'ID: ' + str(id), (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color)
    cv2.imshow('result', src)


while True:
    flag, frame = cap.read()

    # print(flag, frame)
    # cv2.imshow('cap', frame)
    face_detect_demo(frame)

    if ord('q') == cv2.waitKey(10):
        break

cv2.destroyAllWindows()
cap.release()


